
import time
import os
from colorama import Style, Fore
from xml.etree import ElementTree as ET
import glob
import os
import numpy as np
import cv2
import matplotlib.pyplot as plt
from siamrpnpp.utils.bbox import get_axis_aligned_bbox
def overlap_ratio(rect1, rect2):
    '''Compute overlap ratio between two rects
    Args
        rect:2d array of N x [x,y,w,h]
    Return:
        iou
    '''
    # if rect1.ndim==1:
    #     rect1 = rect1[np.newaxis, :]
    # if rect2.ndim==1:
    #     rect2 = rect2[np.newaxis, :]
    #print("iou", len(rect1))
    left = np.maximum(rect1[:,0], rect2[:,0])
    right = np.minimum(rect1[:,0]+rect1[:,2], rect2[:,0]+rect2[:,2])
    top = np.maximum(rect1[:,1], rect2[:,1])
    bottom = np.minimum(rect1[:,1]+rect1[:,3], rect2[:,1]+rect2[:,3])

    #print("iou", len(left),len(right),len(top),len(bottom))
    intersect = np.maximum(0,right - left) * np.maximum(0,bottom - top)

    union = rect1[:,2]*rect1[:,3] + rect2[:,2]*rect2[:,3] - intersect
    iou = intersect / union
    #print("iou", len(iou))
    iou = np.maximum(np.minimum(1, iou), 0)

    return iou

def success_overlap(gt_bb, result_bb, n_frame):
    thresholds_overlap = np.arange(0, 1.05, 0.05)
    success = np.zeros(len(thresholds_overlap))
    iou_ = np.ones(len(gt_bb)) * (-1)
    mask = np.sum(gt_bb > 0, axis=1) == 4
    #print("---",mask,len(gt_bb),len(result_bb),len(mask),len(gt_bb[mask]))
    iou = overlap_ratio(gt_bb[mask], result_bb[mask])
    #print(len(mask), len(iou))
    iou_[mask] = iou
    for i in range(len(thresholds_overlap)):
        success[i] = np.sum(iou > thresholds_overlap[i]) / float(n_frame)
    return success,iou

def success_error(gt_center, result_center, thresholds, n_frame):
    # n_frame = len(gt_center)
    success = np.zeros(len(thresholds))
    dist = np.ones(len(gt_center)) * (-1)
    mask = np.sum(gt_center > 0, axis=1) == 2
    dist[mask] = np.sqrt(np.sum(
        np.power(gt_center[mask] - result_center[mask], 2), axis=1))
    for i in range(len(thresholds)):
        success[i] = np.sum(dist <= thresholds[i]) / float(n_frame)
    return success

def calculate_failures(trajectory):
    """ Calculate number of failures
    Args:
        trajectory: list of bbox
    Returns:
        num_failures: number of failures
        failures: failures point in trajectory, start with 0
    """
    failures = [i for i, x in zip(range(len(trajectory)), trajectory)
            if len(x) == 1 and x[0] == 2]
    num_failures = len(failures)
    return num_failures, failures


def calculate_accuracy(pred_trajectory, gt_trajectory,
                       burnin=0, ignore_unknown=True, bound=None):
    """Caculate accuracy socre as average overlap over the entire sequence
    Args:
        trajectory: list of bbox
        gt_trajectory: list of bbox
        burnin: number of frames that have to be ignored after the failure
        ignore_unknown: ignore frames where the overlap is unknown
        bound: bounding region
    Return:
        acc: average overlap
        overlaps: per frame overlaps
    """
    pred_trajectory_ = pred_trajectory
    if not ignore_unknown:
        unkown = [len(x) == 1 and x[0] == 0 for x in pred_trajectory]

    if burnin > 0:
        pred_trajectory_ = pred_trajectory[:]
        mask = [len(x) == 1 and x[0] == 1 for x in pred_trajectory]
        for i in range(len(mask)):
            if mask[i]:
                for j in range(burnin):
                    if i + j < len(mask):
                        pred_trajectory_[i + j] = [0]
    min_len = min(len(pred_trajectory_), len(gt_trajectory))
    overlaps = region.vot_overlap_traj(pred_trajectory_[:min_len],
                                       gt_trajectory[:min_len], bound)

    if not ignore_unknown:
        overlaps = [x if u else 0 for u in unkown]

    acc = 0
    if len(overlaps) > 0:
        acc = np.nanmean(overlaps)
    return acc, overlaps



def convert_ret(rets):# [xmin, ymin, xmax, ymax] to [x,y,w,h]
    return np.array([rets[:,0],rets[:,1],rets[:,2]-rets[:,0],rets[:,3]-rets[:,1]]).T
def convert_bb_to_center(bboxes):
    return np.array([(bboxes[:, 0] + (bboxes[:, 2] - 1) / 2), (bboxes[:, 1] + (bboxes[:, 3] - 1) / 2)]).T
def convert_bb_to_norm_center(bboxes, gt_wh):
    return convert_bb_to_center(bboxes) / (gt_wh + 1e-16)
def draw_success_precision(success_ret,precision_ret=None, norm_precision_ret=None,axis=[0, 1], xtitle=None, ytitle=None, title=None):
    print(" Success Plot .............................")
    plt.style.use('seaborn-white')
    plt.rcParams['font.family'] = 'Times New Roman'
    plt.rcParams['font.sans-serif'] = ['Times New Roman']
    fontSIZE = 10

    fig, ax = plt.subplots()

    ax.grid(b=True)
    ax.set_aspect(1)
    plt.xlabel('Overlap threshold',fontsize=fontSIZE)
    plt.ylabel('Success rate',fontsize=fontSIZE)
    plt.title('Success plot',fontsize=fontSIZE)
    plt.axis([0, 1] + axis)
    success = {}
    for tracker_name in success_ret.keys():
        value = [v for k, v in success_ret[tracker_name].items()]
        success[tracker_name] = np.mean(value)

    thresholds = np.arange(0, 1.05, 0.05)
    for idx, (tracker_name, auc) in \
            enumerate(sorted(success.items(), key=lambda x: x[1], reverse=True)):
        print(idx, (tracker_name, auc))
        value = [v for k, v in success_ret[tracker_name].items()]
        plt.plot(thresholds, np.mean(value, axis=0), label=tracker_name, linewidth=2)
    ax.legend(loc='lower left', labelspacing=0.2,fontsize=20)
    ax.autoscale(enable=True, axis='both', tight=True)
    xmin, xmax, ymin, ymax = plt.axis()
    ax.autoscale(enable=False)
    ymax += 0.03
    ymin = 0
    plt.axis([xmin, xmax, ymin, ymax])
    plt.xticks(np.arange(xmin, xmax + 0.01, 0.1))
    plt.yticks(np.arange(ymin, ymax, 0.1))
    ax.set_aspect((xmax - xmin) / (ymax - ymin))
    plt.savefig("Success_plot")
    plt.show()


    if precision_ret is None:
        return
    print("Precision Plot .............................")
    fig, ax = plt.subplots()
    ax.grid(b=True)
    ax.set_aspect(50)
    plt.xticks(fontsize=fontSIZE)
    plt.yticks(fontsize=fontSIZE)
    plt.xlabel('Location error threshold',fontsize=fontSIZE)
    plt.ylabel('Precision',fontsize=fontSIZE)

    plt.title('Precision plot',fontsize=fontSIZE)

    plt.axis([0, 50] + axis)
    precision = {}
    thresholds = np.arange(0, 51, 1)
    for tracker_name in precision_ret.keys():
        value = [v for k, v in precision_ret[tracker_name].items()]
        precision[tracker_name] = np.mean(value, axis=0)[20]
    for idx, (tracker_name, pre) in \
            enumerate(sorted(precision.items(), key=lambda x: x[1], reverse=True)):
        print(idx, (tracker_name, pre))
        value = [v for k, v in precision_ret[tracker_name].items()]
        plt.plot(thresholds, np.mean(value, axis=0), label=tracker_name, linewidth=2)
    ax.legend(loc='lower right', labelspacing=0.2,fontsize=20)
    ax.autoscale(enable=True, axis='both', tight=True)
    xmin, xmax, ymin, ymax = plt.axis()
    ax.autoscale(enable=False)
    ymax += 0.03
    ymin = 0
    plt.axis([xmin, xmax, ymin, ymax])
    plt.xticks(np.arange(xmin, xmax + 0.01, 5))
    plt.yticks(np.arange(ymin, ymax, 0.1))
    ax.set_aspect((xmax - xmin) / (ymax - ymin))
    plt.savefig("Precision_plot")
    plt.show()

    # norm precision plot
    if norm_precision_ret:
        print("norm precision plot .............................")
        fig, ax = plt.subplots()
        plt.xticks(fontsize=fontSIZE)
        plt.yticks(fontsize=fontSIZE)
        ax.grid(b=True)
        plt.xlabel('Location error threshold',fontsize=fontSIZE)
        plt.ylabel('Precision',fontsize=fontSIZE)
        plt.title('Normalized Precision plot',fontsize=fontSIZE)
        norm_precision = {}
        thresholds = np.arange(0, 51, 1) / 100
        for tracker_name in precision_ret.keys():
            value = [v for k, v in norm_precision_ret[tracker_name].items()]
            norm_precision[tracker_name] = np.mean(value, axis=0)[20]
        for idx, (tracker_name, pre) in \
                enumerate(sorted(norm_precision.items(), key=lambda x: x[1], reverse=True)):
            print(idx, (tracker_name, pre))
            value = [v for k, v in norm_precision_ret[tracker_name].items()]
            plt.plot(thresholds, np.mean(value, axis=0), label=tracker_name, linewidth=2)
        ax.legend(loc='lower right', labelspacing=0.2,fontsize=10)
        ax.autoscale(enable=True, axis='both', tight=True)
        xmin, xmax, ymin, ymax = plt.axis()
        ax.autoscale(enable=False)
        ymax += 0.03
        ymin = 0
        #print(xmin, xmax, ymin, ymax)
        plt.axis([xmin, xmax, ymin, ymax])
        plt.xticks(np.arange(xmin, xmax + 0.01, 0.1))
        plt.yticks(np.arange(ymin, ymax, 0.1))
        ax.set_aspect((xmax - xmin) / (ymax - ymin))
        plt.savefig("Normalized_Precision_plot")
        plt.show()

def get_plot_img(plot_ret_name=None, plot_pre_ret=None, plot_gt_ret=None):
    #################################################################
    # 一、 Success Plot

    success_ret = {}
    for i, ret_name in enumerate(plot_ret_name):
        success_ret_ = {}
        #print(len(plot_gt_ret[i]))
        for j in range(len(plot_gt_ret[i])):
            gt_traj = plot_gt_ret[i][j]

            n_frame = len(gt_traj)
            tracker_traj = plot_pre_ret[i][j]
            #print(gt_traj,tracker_traj,n_frame)
            success_ret_[str(j)],iou = success_overlap(gt_traj, tracker_traj, n_frame)
            #np.save("overlap\\{}_{}.npy".format(ret_name,j),iou)
        success_ret[ret_name] = success_ret_
    # 二、 Precision Plot
    precision_ret = {}
    for i, ret_name in enumerate(plot_ret_name):
        precision_ret_ = {}
        for j in range(len(plot_gt_ret[i])):
            gt_traj = plot_gt_ret[i][j]
            n_frame = len(gt_traj)
            tracker_traj = plot_pre_ret[i][j]
            gt_center = convert_bb_to_center(gt_traj)
            # print(gt_center)
            tracker_center = convert_bb_to_center(tracker_traj)
            # print(tracker_center)
            thresholds = np.arange(0, 51, 1)
            precision_ret_[str(j)] = success_error(gt_center, tracker_center, thresholds, n_frame)
        precision_ret[ret_name] = precision_ret_
    # 三、Norm Precision Plot
    norm_precision_ret = {}
    for i, ret_name in enumerate(plot_ret_name):
        norm_precision_ret_ = {}
        for j in range(len(plot_gt_ret[i])):
            gt_traj = plot_gt_ret[i][j]
            n_frame = len(gt_traj)
            tracker_traj = plot_pre_ret[i][j]
            gt_center_norm = convert_bb_to_norm_center(gt_traj, gt_traj[:, 2:4])
            tracker_center_norm = convert_bb_to_norm_center(tracker_traj, gt_traj[:, 2:4])
            # print(gt_center_norm)
            # print(tracker_center_norm)
            thresholds = np.arange(0, 51, 1) / 100
            norm_precision_ret_[str(j)] = success_error(gt_center_norm, tracker_center_norm, thresholds, n_frame)
        norm_precision_ret[ret_name] = norm_precision_ret_
    draw_success_precision(success_ret=success_ret, precision_ret=precision_ret, norm_precision_ret=norm_precision_ret,
                           xtitle="Overlap threshold", ytitle="Precision", title="Precision plots of OPE")



def from_txt(gt_path):
    gts = []
    with open(gt_path, 'r') as f:
        while True:
            line = f.readline()
            if line == '':
                gts = np.array(gts, dtype=np.float32)
                return gts
            if ',' in line:
                gt_pos = line.split(',')
            else:
                gt_pos = line.split()
            gt_pos_int = [(float(element)) for element in gt_pos]
            gts.append(gt_pos_int)

def from_gt_txt(gt_path):
    gts = []
    with open(gt_path, 'r') as f:
        while True:
            line = f.readline()
            if line == '':
                gts = np.array(gts, dtype=np.float32)
                return gts
            if ',' in line:
                gt_pos = line.split(',')
            else:
                gt_pos = line.split()
            gt_pos_int = [(float(element)) for element in gt_pos]
            cx, cy, w, h = get_axis_aligned_bbox(np.array(gt_pos_int))  # [x, y, w, h] --> [cx, cy, w, h]
            gt_bbox_ = [cx - (w - 1) / 2, cy - (h - 1) / 2, w, h]  # [topx,topy,w,h]
            gts.append(gt_bbox_)
def get_predict_box():

    pre_path = os.path.join("results", "train", "SiamRPN++")
    pre_labels = os.listdir(pre_path)

    gt_root_path = os.path.join("datasets", "trainval", "trainval")

    #
    all_list = []
    for prename in pre_labels:
        if prename == "img":
            continue
        #print(prename)

        pr_path = os.path.join(pre_path,prename)
        pre_box = from_txt(pr_path)
        #
        name = prename.split(".")[0]
        gt_path = os.path.join(gt_root_path,name,"groundtruth.txt")
        gt_box = from_gt_txt(gt_path)
        print(len(pre_box), len(gt_box))
        assert len(pre_box)==len(gt_box)
        all_list.append(["tracker", pre_box, gt_box])




    print("*" * 40)
    ##############################################
    plot_ret_name = []
    plot_pre_ret = []
    plot_gt_ret = []
    for box_info in all_list:
        if box_info[0] in plot_ret_name:
            index = plot_ret_name.index(box_info[0])
            plot_pre_ret[index].append(box_info[1])
            plot_gt_ret[index].append(box_info[2])
        else:
            plot_ret_name.append(box_info[0])
            plot_pre_ret.append([box_info[1]])
            plot_gt_ret.append([box_info[2]])
    assert len(plot_ret_name) == len(plot_pre_ret), ("{}预测长度：{} 不等于 标签长度{}".format(plot_ret_name, len(plot_ret_name), len(plot_pre_ret)))
    return plot_ret_name, plot_pre_ret, plot_gt_ret



if __name__ == '__main__':
    print("start.....")
    plot_ret_name, plot_pre_ret, plot_gt_ret = get_predict_box()

    get_plot_img(plot_ret_name, plot_pre_ret, plot_gt_ret)







